import time

import requests
import os
from service_conf import SERVER_URL

# 定义发送POST请求的函数
def send_post_request(url, data):
    response = requests.post(url, json=data)
    return response.json()


# 定义处理RVC任务的函数
def run_rvc_task(task_id, user_id, song_id):
    # 构建请求数据
    data = {
        'task_id': task_id,
        'user_id': user_id or '__defualt_user_id',  # 如果没有提供user_id，使用默认值
        'song_id': song_id or '__default_song_id'  # 如果没有提供song_id，使用默认值
    }

    # 发送POST请求到服务器的/run_rvc路由
    result = send_post_request(f"{SERVER_URL}/run_rvc", data)

    # 检查响应状态
    if result['status'] == 'vc_runing':
        print(
            f"任务ID {result['task_id']}已成功启动，AI声音文件路径：{result['ai_vocal']}，歌曲文件路径：{result['song']}。")
        return result['ai_vocal'],result['song']
    else:
        print(f"任务ID {result['task_id']}启动失败，错误信息：{result.get('error', '未知错误')}")
    return None,None


# 定义下载预测结果的函数
def download_prediction(song_name,output_dir):
    prediction_url = f"{SERVER_URL}/get_prediction/{output_dir}"
    response = requests.get(prediction_url, stream=True)

    # 检查响应状态码
    if response.status_code == 200:
        # 将文件保存到本地
        print(prediction_url)
        file_name = os.path.basename(prediction_url).split('/')[-1]
        saved_file_name = f'{song_name}.wav'
        with open(saved_file_name, 'wb') as file:
            for chunk in response.iter_content(chunk_size=1024):
                if chunk:  # 过滤掉保持连接的新块
                    file.write(chunk)
        print(f"预测文件 '{file_name}' 下载成功。")
    elif response.status_code == 404:
        print("错误：未找到预测结果。")
    else:
        print(f"错误：下载预测文件失败，状态码：{response.status_code}")


def start_gen_and_download(task_id,user_id,song_id):
    """
    1发起AI 转换音色命令，
    2等生成完毕下载文件
    :param task_id: 输入trainer 返回的task id
    :param user_id: 随便输入
    :param song_id: 输入以下歌曲之一 当你老了/盛夏的果实/如果没有你/我要找到你/小美满/这世界那么多人
    :return:
    """
    # 模型名字， 也是train 的task id

    # 运行RVC任务
    print(' # 运行RVC任务')
    vocal,song = run_rvc_task(task_id, user_id, song_id)
    if song:
        # 要等待10多秒生成时间
        time.sleep(20)
        print(f'歌曲地址为{song}')
        download_prediction(song_id,os.path.dirname(song))
    else:
        print('没有返回歌曲地址')

if __name__ == "__main__":
    #task_id = 'training_user123_20240405153939'
    task_id = 'girl'
    user_id = 'yuhang'
    song_id = '最美不过是朝晖'  # sx_d_gs rgmyn zsj_nm_dr xmm wyzdn
    # 当你老了/盛夏的果实/如果没有你/我要找到你/小美满/这世界那么多人
    start_gen_and_download(task_id, user_id, song_id)

